import scipy.io as scio
import matplotlib.pyplot as plt
from scipy.misc import imread,imshow
import tensorflow as tf
import numpy as np

data=scio.loadmat("/home/chen/DCNN_caffe-Tensorflow/test1.mat")


def weight_variable(name,shape):
    initial=data[name]
    x = tf.sqrt(tf.cast(tf.shape(initial)[1], tf.float32))
    x=tf.cast(x,tf.int32)
    y = tf.shape(initial)[2]
    initial = tf.reshape(initial, [-1,x,x,y])
    initial = tf.transpose(initial, [2,1,0,3])       #[hei,wid,batch,channel]
    initial=tf.cast(tf.reshape(initial,shape),tf.float32)
    return tf.Variable(initial)


def bias_variable(name,shape):
    #initial=tf.constant(data[name],shape)
     initial=data[name]
     initial=tf.cast(tf.reshape(initial,shape),tf.float32)
     return tf.Variable(initial)


def conv2d(x,W):
    return tf.nn.conv2d(x,W,strides=[1,1,1,1],padding='SAME')
#[1,x_movement,y_movement,1], padding=SAME/VALID


def max_pool_2x2(x):
    return tf.nn.max_pool(x,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
#[1,x_movement,y_movement,1], padding=SAME/VALID





def getImage(image_path,knn_path,closedform_path):
    ClosedForm_image=tf.cast(imread(closedform_path),tf.float32)
    KNN_image=tf.cast(imread(knn_path),tf.float32)
    x_image=tf.cast(imread(image_path),tf.float32)
    with tf.Session() as sess:
         s=x_image[:,:,0]*x_image[:,:,0]+x_image[:,:,1]*x_image[:,:,1]+x_image[:,:,2]*x_image[:,:,2]
         s=sess.run(s)
    s=tf.pack(s)
    x_image_r= tf.div(x_image[:,:,0],s)
    x_image_g= tf.div(x_image[:,:,1],s)
    x_image_b= tf.div(x_image[:,:,2],s)
    with tf.Session() as sess:
         ClosedForm_image=sess.run(ClosedForm_image)
         KNN_image=sess.run(KNN_image)
         # print(np.shape(sess.run(x_image_r)))
         #x_image=[sess.run(x_image_r),sess.run(x_image_g),sess.run(x_image_b)]
         x_image=[sess.run(x_image_r),sess.run(x_image_g),sess.run(x_image_b),ClosedForm_image, KNN_image]
         x_image = tf.pack(x_image,axis=2)
         x_image=sess.run(x_image)
         x=tf.shape(x_image)[0]
         y=tf.shape(x_image)[1]
         c=tf.shape(x_image)[2]
         x_image=tf.reshape(x_image,[1,x,y,c])
         return x_image



#x_image=tf.reshape(xs,[-1,28,28,1])    #[n_samples,width,height,channel]
image_path='/home/chen/DCNN_caffe-Tensorflow/image.png'
closedform_path='/home/chen/DCNN_caffe-Tensorflow/ClosedForm.png'
knn_path='/home/chen/DCNN_caffe-Tensorflow/KNN.png'
x_image=getImage(image_path,closedform_path,knn_path)
x_image=x_image-0.5
trimap=imread('/home/chen/DCNN_caffe-Tensorflow/trimap.png')



#================================================
#conv1 layer
W_conv1 = weight_variable('weights_conv1',[9,9,5,64])    #9x9 patch, in size 5, out size 64
#W_conv1=W_conv1[:,:,0:3,:]
b_conv1=bias_variable('biases_conv1',[64])
h_conv1=tf.nn.relu(conv2d(x_image,W_conv1)+b_conv1) #output size 
#h_conv1=conv2d(x_image,W_conv1)+b_conv1
h_conv1=tf.nn.dropout(h_conv1,0.7)



#conv2_1 layer
W_conv2_1 = weight_variable('weights_conv2_1',[1,1,64,64])    #1x1 patch, in size 64, out size 64
b_conv2_1=bias_variable('biases_conv2_1',[64])
h_conv2_1=tf.nn.relu(conv2d(h_conv1,W_conv2_1)+b_conv2_1) 


#conv2_2 layer
W_conv2_2 = weight_variable('weights_conv2_2',[1,1,64,64])    #1x1 patch, in size 64, out size 64
b_conv2_2=bias_variable('biases_conv2_2',[64])
h_conv2_2=tf.nn.relu(conv2d(h_conv2_1,W_conv2_2)+b_conv2_2) 

 
#conv2_3 layer
W_conv2_3 = weight_variable('weights_conv2_3',[1,1,64,64])    #1x1 patch, in size 64, out size 64
b_conv2_3=bias_variable('biases_conv2_3',[64])
h_conv2_3=tf.nn.relu(conv2d(h_conv2_2,W_conv2_3)+b_conv2_3) 



#conv2_4 layer
W_conv2_4 = weight_variable('weights_conv2_4',[1,1,64,32])    #1x1 patch, in size 64, out size 32
b_conv2_4=bias_variable('biases_conv2_4',[32])
h_conv2_4=tf.nn.relu(conv2d(h_conv2_3,W_conv2_4)+b_conv2_4) 



#conv1 layer
W_conv3 = tf.cast(data['weights_conv3'],tf.float32)    #5x5 patch, in size 32, out size 1
W_conv3=tf.reshape(W_conv3,[5,5,32,1])
b_conv3=bias_variable('biases_conv3',[1])
h_conv3=conv2d(h_conv2_4,W_conv3)+b_conv3 #output size 
#h_conv1=conv2d(x_image,W_conv1)+b_conv1





prediction=h_conv3+0.5
#prediction=prediction*255
with tf.Session() as sess:
     sess.run(tf.initialize_all_variables())
     x=tf.shape(prediction)[1]
     y=tf.shape(prediction)[2]
     c=tf.shape(prediction)[3]
     prediction=tf.reshape(prediction,[x,y])
     #print(sess.run(prediction))
     #prediction=tf.pack(prediction,2)
     #prediction=tf.cast(prediction,tf.uint8)
     prediction=sess.run(prediction)    
     #prediction_im=Image.fromarray(prediction)
     #prediction=tf.unpack(prediction)
     #prediction_im.save('/home/temp.png')







prediction=np.array(prediction,np.float32)
#prediction=np.dot(prediction,255)


prediction[prediction<0]=0
prediction[prediction>255]=255
x=np.shape(prediction)[0]
y=np.shape(prediction)[1]


for i in range (0,x-1):
    for j in range (0,y-1):
        if (trimap[i,j].max()<1):
            prediction[i,j]=0.0
        if (trimap[i,j].min()>254):
            prediction[i,j]=255.0


closedform=imread(closedform_path)
knn=imread(knn_path)
for i in range (0,x-1):
    for j in range (0,y-1):
        if (np.abs(closedform[i,j]-knn[i,j])<30):
            prediction[i,j]=closedform[i,j]
imshow(prediction)





